import torch 
from torch import nn
from torch.nn import functional as F
import pickle
from typing import NamedTuple
import numpy as np
from scipy import sparse
from .utils import compute_v_normals, compute_normals, batch_rodrigues
from ..vhap.model.flame import FlameHead

FLAME_MODEL_PATH = "asset/flame/flame2023.pkl"
FLAME_MESH_PATH = "asset/flame/head_template_mesh.obj"
FLAME_PARTS_PATH = "asset/flame/FLAME_masks.pkl"
FLAME_LMK_PATH = "asset/flame/landmark_embedding_with_eyes.npy"
FLAME_TEX_PATH = "asset/flame/FLAME_texture.npz"
FLAME_PAINTED_TEX_PATH = "asset/flame/tex_mean_painted.png"
FLAME_UVMASK_PATH = "asset/flame/uv_masks.npz"

# %% Flame
class StructFlame(NamedTuple):
    f: np.ndarray
    v_template: np.ndarray
    kintree_table: np.ndarray
    weights: np.ndarray
    posedirs: np.ndarray
    shapedirs: np.ndarray
    J: np.ndarray
    J_regressor: sparse.csc_matrix
    bs_style: str
    bs_type: str
    

# %% LBS
def transform_mat(R, t):
    """Creates a batch of transformation matrices
    Args:
        - R: Bx3x3 array of a batch of rotation matrices
        - t: Bx3x1 array of a batch of translation vectors
    Returns:
        - T: Bx4x4 Transformation matrix
    """
    # No padding left or right, only add an extra row
    return torch.cat([F.pad(R, [0, 0, 0, 1]), F.pad(t, [0, 0, 0, 1], value=1)], dim=2)

def vertices2joints(J_regressor, vertices):
  """
  根据J_regressor(shape为BxJxV)的权重分配,获取BxJx3的关节点坐标
  vertices: BxVx3
  """
  return torch.einsum("bik,ji->bjk", [vertices, J_regressor])

def blend_shapes(betas, shape_disps):
  """
  betas: Bx(num_betas) Bx[300+100] shape:300 expr:100
  """
  blend_shape = torch.einsum("bl,mkl->bmk", [betas, shape_disps])
  return blend_shape

def batch_rigid_transform(rot_mats, joints, parents, dtype=torch.float32):
    """
    Applies a batch of rigid transformations to the joints

    Parameters
    ----------
    rot_mats : torch.tensor BxNx3x3
        Tensor of rotation matrices
    joints : torch.tensor BxNx3
        Locations of joints
    parents : torch.tensor BxN
        The kinematic tree of each object
    dtype : torch.dtype, optional:
        The data type of the created tensors, the default is torch.float32

    Returns
    -------
    posed_joints : torch.tensor BxNx3
        The locations of the joints after applying the pose rotations
    rel_transforms : torch.tensor BxNx4x4
        The relative (with respect to the root joint) rigid transformations
        for all the joints
    """

    joints = torch.unsqueeze(joints, dim=-1)

    rel_joints = joints.clone().contiguous()
    rel_joints[:, 1:] = rel_joints[:, 1:] - joints[:, parents[1:]]

    transforms_mat = transform_mat(rot_mats.view(-1, 3, 3), rel_joints.view(-1, 3, 1))
    transforms_mat = transforms_mat.view(-1, joints.shape[1], 4, 4)

    transform_chain = [transforms_mat[:, 0]]
    for i in range(1, parents.shape[0]):
        # Subtract the joint location at the rest pose
        # No need for rotation, since it's identity when at rest
        curr_res = torch.matmul(transform_chain[parents[i]], transforms_mat[:, i])
        transform_chain.append(curr_res)

    transforms = torch.stack(transform_chain, dim=1)

    # The last column of the transformations contains the posed joints
    posed_joints = transforms[:, :, :3, 3]

    joints_homogen = F.pad(joints, [0, 0, 0, 1])

    rel_transforms = transforms - F.pad(
        torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0]
    )

    return posed_joints, rel_transforms

def lbs(
    pose,
    v_shaped,
    posedirs,
    J_regressor,
    parents,
    lbs_weights,
    pose2rot=True,
    dtype=torch.float32,
):
    batch_size = pose.shape[0]
    device = pose.device

    # Get the joints
    # NxJx3 array
    J = vertices2joints(J_regressor, v_shaped)

    # 3. Add pose blend shapes
    # N x J x 3 x 3
    ident = torch.eye(3, dtype=dtype, device=device)
    if pose2rot:
        rot_mats = batch_rodrigues(pose.view(-1, 3), dtype=dtype).view(
            [batch_size, -1, 3, 3]
        )

        pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1]) # 因为减去了单位矩阵,所以是相对旋转,表示坐标点的相对位移
        # (N x P) x (P, V * 3) -> N x V x 3 , P是4x9, posedirs是36x5023x3, 表示每个点被关节旋转的权重
        pose_offsets = torch.matmul(pose_feature, posedirs).view(batch_size, -1, 3)
    else:
        pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
        rot_mats = pose.view(batch_size, -1, 3, 3)

        pose_offsets = torch.matmul(pose_feature.view(batch_size, -1), posedirs).view(
            batch_size, -1, 3
        )

    v_posed = pose_offsets + v_shaped

    # 4. Get the global joint location
    J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)

    # 5. Do skinning:
    # W is N x V x (J + 1)
    W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
    # (N x V x (J + 1)) x (N x (J + 1) x 16)
    num_joints = J_regressor.shape[0]
    T = torch.matmul(W, A.view(batch_size, num_joints, 16)).view(batch_size, -1, 4, 4)

    homogen_coord = torch.ones(
        [batch_size, v_posed.shape[1], 1], dtype=dtype, device=device
    )
    v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
    v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))

    verts = v_homo[:, :, :3, 0]

    return verts, T, J_transformed, A

# %% Flame LBS
class FlameLBS(nn.Module):
  dtype = torch.float32
  
  def register_from_numpy(self, param, dtype=None):
      dtype = dtype or self.dtype
      return nn.Parameter(torch.from_numpy(param).to(dtype), requires_grad=False)
  
  def __init__(self, 
               flame_model_path = FLAME_MODEL_PATH,
               zero_centered_at_root_node=False,
               ):
    super().__init__()
    with open(flame_model_path, "rb") as f:
        ss = pickle.load(f, encoding="latin1")
        flame_model = StructFlame(**ss)
    # 如果使用VHAP flame且添加teeth, 则顶点数为5143
    self.faces = self.register_from_numpy(flame_model.f) 
    self.v_template = self.register_from_numpy(flame_model.v_template)               # [5023,3]
    self.J_regressor = self.register_from_numpy(flame_model.J_regressor.todense())   # [5, 5023] 关节位置权重, J_regressor.sum(dim=1) == 1, 基本每个关节位置相关顶点数不多于15个
    self.lbs_weights = self.register_from_numpy(flame_model.weights)                 # [5023, 5] 顶点lbs-关节权重, lbs_weights.sum(dim=1) == 1, 表示每个顶点受到的关节影响权重
    self.shapedirs = self.register_from_numpy(np.array(flame_model.shapedirs))       # [5023, 3, 400] 400 = 300(shape) + 100(expression)
    self.posedirs = self.register_from_numpy(flame_model.posedirs.reshape(-1, 36).T) # pose blend shapes [36, 15429] <= [4*9, 5023*3] 4个关节&旋转,每个旋转3x3的旋转矩阵 <= [5023, 3, 36] flame初始格式
    self.parents = self.register_from_numpy(flame_model.kintree_table[0], torch.long)     # kinematic tree [5] 运动链 -1 根（头部） 0 颈 1 下颚 1 左眼 1 右眼
    self.parents.data[0] = -1
    self.zero_centered_at_root_node = zero_centered_at_root_node # root节点置空间中央
    self._v_static = None
    self._n_static = compute_v_normals(self.v_static[None], self.faces)[0]
    self._f_normal = compute_normals(self.v_static[None], self.faces)[0]
    
  @property
  def vertices_num(self):
    return self.v_template.size(0)
  
  @property
  def faces_num(self):
    return self.faces.size(0)
    
  @property
  def face_normals(self):
    return self._f_normal
    
  def load_from_flame(self, 
                      flame: FlameHead
                      ):
    """
    Load from FlameHead
    """
    self.v_template = nn.Parameter(flame.v_template, requires_grad=False)
    self.shapedirs = nn.Parameter(flame.shapedirs, requires_grad=False)
    self.posedirs = nn.Parameter(flame.posedirs, requires_grad=False)
    self.J_regressor = nn.Parameter(flame.J_regressor, requires_grad=False)
    self.parents = nn.Parameter(flame.parents, requires_grad=False)
    self.lbs_weights = nn.Parameter(flame.lbs_weights, requires_grad=False)
    self.faces = nn.Parameter(flame.faces, requires_grad=False)
    self._n_static = compute_v_normals(self.v_static[None], self.faces)[0]
    self._f_normal = compute_normals(self.v_static[None], self.faces)[0]
    
  @property
  def v_static(self):
    if self._v_static is None:
      return self.v_template
    else:
      return self._v_static
    
  @property
  def n_static(self):
    return self._n_static
    
  def set_v_static(self, 
                     shape: torch.Tensor, # (300,)
                     static_offset: torch.Tensor = None,
                     ):
    betas = torch.cat([shape, torch.zeros(100).to(shape)], dim=0).unsqueeze(0)
    template_vertices = self.v_template.unsqueeze(0)    
    v_shaped = template_vertices + blend_shapes(betas, self.shapedirs)
    if static_offset is not None:
        v_shaped += static_offset
    self._v_static = v_shaped[0].detach().clone()
    
  def set_n_static(self, tri: torch.Tensor):
    self._n_static = compute_v_normals(self.v_static[None], tri)[0]
    
  def forward(
        self,
        shape,
        expr,
        rotation,
        neck,
        jaw,
        eyes,
        translation=None,
        static_offset=None,
    ):
        """
        Input:
            shape_params: N X number of shape parameters
            expression_params: N X number of expression parameters
            pose_params: N X number of pose parameters (6)
        return:d
            vertices: N X V X 3
            landmarks: N X number of landmarks X 3
        """
        batch_size = shape.shape[0]

        betas = torch.cat([shape, expr], dim=1)
        full_pose = torch.cat([rotation, neck, jaw, eyes], dim=1)
        template_vertices = self.v_template.unsqueeze(0).expand(batch_size, -1, -1)
        
        # Add shape contribution
        v_shaped = template_vertices + blend_shapes(betas, self.shapedirs)

        # Add personal offsets
        if static_offset is not None:
            v_shaped += static_offset

        vertices, v_transmat, joints, j_transmat = lbs(
            full_pose,
            v_shaped,
            self.posedirs,
            self.J_regressor,
            self.parents,
            self.lbs_weights,
            dtype=self.dtype,
        )

        if self.zero_centered_at_root_node:
            vertices = vertices - joints[:, [0]]
            joints = joints - joints[:, [0]]

        if not translation is None:
            vertices = vertices + translation[:, None, :]
            joints = joints + translation[:, None, :]
            
        return vertices, v_transmat, joints    